Stabilize DeepCompile ZeRO-3 memory scheduling - #8169
Conversation
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
tohtana
changed the title
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 55779855e3
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
| m.def("register_graph_z3", | ||
| &dc::register_graph_z3, | ||
| "Register graph with a list of ds parameter ids"); | ||
| m.def("set_z3_gather_buffer_pool_budget_for_test", |
There was a problem hiding this comment.
wondering if these new items are just test purpose only?
There was a problem hiding this comment.
Yes, these three new APIs are for testing purpose only. We need them to set or inspect the internal state of the memory pool. It doesn't look great to have these as public APIs, but I didn't have an idea about other approaches.
There was a problem hiding this comment.
just curious if it is possible to do integration test rather than unit test on this level? For example, instead of writing the scaffold to test GatherBufferPool, can do a test of functions that call the natural creation of this class, e.g. calling from Z3CustomOpExecutor?
There was a problem hiding this comment.
Yes, that is actually how the current tests are structured. They do not construct or call GatherBufferPool directly. register_graph_z3() creates Z3CustomOpExecutor with the shared pool, and the CUDA test runs the real allgather_param, wait_allgather, and release_param ops.
The _for_test APIs are limited to making the budget and allocator-pressure inputs deterministic and exposing accounting of internal values.
There was a problem hiding this comment.
OK, your last sentence explains my confusion. OK, maybe we can call it deterministic or pre determined size etc?
| if (!enabled_ || pressure_recovery_in_progress_) { return at::Tensor(); } | ||
|
|
||
| Entry* best = nullptr; | ||
| for (auto& entry : entries_) { |
There was a problem hiding this comment.
is entries_ ordered? if not, is the full loop expensive vs an max heap like solution?
There was a problem hiding this comment.
entries_ is currently an unsorted std::vector as we don't simply choose the maximum one (best fit to capacity, LRU for eviction, search by storage identity, etc.)
However, I don't think the scan is significant overhead. I expect the pool to contain tens of entries rather than thousands because it is byte-bounded and the dominant gather buffers are relatively large. Given an 80GB GPU, the hard cap is 80 GiB / 32 = 2.5 GiB. If we the average buffer size is 64MB, we will have 80 entries. This should be negligible in the C++ code.
|
general question, do we have experiments results showcasing the changes' benefits? |
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
| pool = std::make_shared<GatherBufferPool>(); | ||
| weak_gather_buffer_pool = pool; | ||
| if (gather_buffer_pool_test_budget) { | ||
| pool->setBudgetForTest(gather_buffer_pool_test_budget.value()); |
There was a problem hiding this comment.
setBudgetForTest is only called when it is test purpose?
There was a problem hiding this comment.
Yes. gather_buffer_pool_test_budget is unset by default and is populated only by set_z3_gather_buffer_pool_budget_for_test().
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Yes. We have some measurements although I have not rerun the exact current head yet. The experiment used
The results are:
|
That's great, maybe add that to the PR's description? |
|
the gather buffer pool and the memory budget seem to be two separate optimization, worth considering splitting to to PRs? |
3 participants
Depends on #8159.
DeepCompile's ZeRO-3 scheduler does not consistently account for memory pressure across gather/release ordering, prefetch, selective gathering, and native gather-buffer reuse. This can retain excess gathered storage or trigger a full-parameter gather while Dynamo evaluates guards.
This PR adds rank-consistent scheduler budgeting and diagnostics, coordinates prefetch and selective gathering with graph profiling, reuses native gather storage within a bounded pressure-aware lifecycle, and avoids guard-time full-parameter gathers while the DeepCompile eager fallback is active.